#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :huaban_pins.py
# @Time      :2023/7/27 
# @Author    :CL
# @email     :1037654919@qq.com

import os
import json
import random
import time

import imageio
from tqdm import tqdm
from utils import MongoDBUtil, get_html, mongo_manager
from multiprocessing import Pool
import requests
from huaban_util import Huaban

MongoDBUtil = MongoDBUtil()
huaban_boards = mongo_manager("huaban_boards", db="car_images")
huaban = Huaban()
images_keywords = mongo_manager("images_keywords", db="car_images")
proxies = {'http': '127.0.0.1:15732',
           'https': '127.0.0.1:15732'}


def get_proxy():
    port = random.randint(24000, 24400)
    proxies = {'http': f'http://zheng123:zheng123@haproxy.iinti.cn:{port}',
               'https': f'http://zheng123:zheng123@haproxy.iinti.cn:{port}'}
    return proxies


def insert_data(tablename, datas):
    if len(datas) > 0 and datas is not None:
        for data in datas:
            # data['_id'] = data['id']
            try:
                # print(data)
                MongoDBUtil.insert_one(tablename, data)
            except Exception as e:
                print(e)
    else:
        pass


def down_image(file, url):  # 下载图片方法
    print("开始下载：", url)
    for i in range(5):
        try:
            response = requests.get(url, proxies=proxies, timeout=10)
            with open(file, 'wb') as fd:
                fd.write(response.content)
            requests.session().close()
            return 1
        except Exception as e:
            time.sleep(5)
    print("下载失败了", url)
    if os.path.exists(file):
        os.remove(file)
    return 0


def parse_url(seed):
    huaban_boards = mongo_manager("huaban_boards", db="car_images")
    keyword = str(seed['keyword_e']).strip()
    board_id = seed['_id']
    pin_count = seed['pin_count']
    url = "https://api.huaban.com/boards/{}/pins".format(board_id)
    print('begin 访问网页:', url, '  pin_count:', pin_count)
    path = f'/media/chenglei3/77D014CED257D1411/images/huaban/{keyword}/'
    os.makedirs(path, exist_ok=True)

    response = huaban.get_pins(board_id, maxs=None)
    board = response['board']
    desc = board['description']
    category_id = board['category_id']
    last_pin_id = None
    success_count = 0
    flag = 40
    while flag >= 40:
        try:
            response = huaban.get_pins(board_id, maxs=last_pin_id)
            if response:
                pins = response['pins']
                flag = len(pins)
                last_pin_id = pins[-1]['pin_id']
                for pin in pins:
                    # print(pin['pin_id'])
                    datas = {'tags': pin['tags']}
                    tag_list = pin['tags']
                    datas['_id'] = pin['pin_id']
                    datas['pin_id'] = pin['pin_id']
                    datas['keyword'] = seed['keyword']
                    datas['board_id'] = pin['board_id']
                    datas['raw_text'] = pin['raw_text']
                    datas['file'] = pin['file']
                    href = 'https://huaban.com/pins/' + str(pin['pin_id'])
                    datas['href'] = href
                    src = 'https://gd-hbimg.huaban.com/' + str(pin['file']['key'])
                    datas['src'] = src
                    try:
                        MongoDBUtil.insert_one('huaban_pins', datas)
                    except Exception as e:
                        print(e)
                        # continue
                    name = src.split('/')[-1]
                    # 保存图片
                    result = down_image(file=f'{path}{name}', url=src)
                    # 保存文本
                    if result == 1:
                        description = str(pin['raw_text']).rsplit('|', 1)[0]
                        if str(description).strip() == '':
                            description = desc
                        image_path = f'{path}{name}'
                        img = imageio.v2.imread(image_path)
                        # 获取图片尺寸
                        height, width = img.shape[0], img.shape[1]
                        json_data = {'_id': pin['pin_id'], 'image_url': src, 'size': f'{height},{width}',
                                     'label': ','.join(tag_list), 'desc': description}
                        # json_name = filename.split('.')[0]
                        with open(f'{path}{name.split(".")[0]}.json', 'w') as f:
                            json.dump(json_data, f, ensure_ascii=False)
                        success_count += 1
                        datas['status'] = 'success'
                        huaban_pins = mongo_manager("huaban_pins", db="car_images")
                        huaban_pins.updateOne({'_id': datas['_id']}, datas)
                        huaban_pins.close()
        except Exception as e:
            print(e)
        if len(os.listdir(path)) > 2000:
            print(f'{keyword} enough')
            huaban_boards.updateMany({'keyword_e': seed['keyword_e']}, {'status': 'enough'})
            break
    seed['success_count'] = success_count
    if success_count > 0:
        seed['description'] = desc
        seed['category_id'] = category_id
        seed['status'] = 'success'
    else:
        seed['status'] = 'unsuccess'
    huaban_boards.updateOne({'_id': seed['_id']}, seed)

    huaban_boards.close()
    return


def run_keyword(seed):
    huaban_boards = mongo_manager("huaban_boards", db="car_images")

    keyword_en = seed['keyword_en'].strip()
    path = f'/media/chenglei3/77D014CED257D1411/images/huaban/{keyword_en}/'
    os.makedirs(path, exist_ok=True)
    print('keyword_en:', keyword_en, len(os.listdir(path)))
    if len(os.listdir(path)) >= 2000:
        huaban_boards.updateMany({"status": None, "keyword_e": seed['keyword_en']}, {"status": "enough"})
    else:
        seeds = huaban_boards.findAll({"status": None, "keyword_e": seed['keyword_en']})
        for seed in seeds:
            if len(os.listdir(path)) >= 2000:
                print(seed['keyword_e'], ' running enough')
                huaban_boards.updateMany({"status": None, "keyword_e": seed['keyword_e']}, {"status": "enough"})
                break
            else:
                parse_url(seed)
    huaban_boards.close()

def main():
    lists = []
    seeds = images_keywords.findAll({'huaban': 'done'})
    for seed in seeds:
        lists.append(seed)
    pool = Pool(processes=10)
    pool.map(run_keyword, lists)
    pool.close()  # 关闭进程池，不再接受新的进程
    pool.join()


if __name__ == "__main__":
    main()
    huaban_boards.close()
    images_keywords.close()
